# 按钮-生成树
import os

from graphviz import Digraph

from analyze import python_parser, format_tree, get_truth_trees, print_tree
from analyzeJava import format_java_tree, java_parser
from analyze_cpp import format_tree_cpp, print_tree_cpp, cpp_parser
from entity import If_node


# 将树可视化并保存为图像
def visualize_tree(root: If_node, dot=None):
    if dot is None:
        dot = Digraph(comment='AST Tree')

    # 为每个节点分配一个唯一的ID
    node_id = id(root)
    label = repr(root)
    dot.node(str(node_id), label)

    # 遍历子节点
    for child in root.children:
        child_id = id(child)
        dot.node(str(child_id), repr(child))
        dot.edge(str(node_id), str(child_id))
        # 递归调用处理子节点
        visualize_tree(child, dot)

    return dot


# 保存图像
def save_ast_image(if_tree: If_node, output_dir: str):
    # 确保输出目录存在
    if not os.path.exists(output_dir):
        os.makedirs(output_dir)

    # 设置图像文件名和路径
    output_path = os.path.join(output_dir, 'ast_tree')

    # 创建并生成图像
    dot = visualize_tree(if_tree)
    dot.format = 'png'  # 图像格式可以是 'png', 'pdf', 'jpeg', 等
    dot.attr(dpi='300')  # 设置图片分辨率（可选）
    dot.render(output_path)  # 保存图片到指定路径

def save_ast_image_cpp(if_tree: If_node, output_dir: str):
    # 确保输出目录存在
    if not os.path.exists(output_dir):
        os.makedirs(output_dir)

    # 设置图像文件名和路径
    output_path = os.path.join(output_dir, 'ast_tree_cpp')

    # 创建并生成图像
    dot = visualize_tree(if_tree)
    dot.format = 'png'  # 图像格式可以是 'png', 'pdf', 'jpeg', 等
    dot.attr(dpi='300')  # 设置图片分辨率（可选）
    dot.render(output_path)  # 保存图片到指定路径


def save_ast_image_java(if_tree: If_node, output_dir: str):
    # 确保输出目录存在
    if not os.path.exists(output_dir):
        os.makedirs(output_dir)

    # 设置图像文件名和路径
    output_path = os.path.join(output_dir, 'ast_tree_java')

    # 创建并生成图像
    dot = visualize_tree(if_tree)
    dot.format = 'png'  # 图像格式可以是 'png', 'pdf', 'jpeg', 等
    dot.attr(dpi='300')  # 设置图片分辨率（可选）
    dot.render(output_path)  # 保存图片到指定路径


# 按钮-生成语法树
def generateTree(code: str):
    # ----------------------生成中间表示（抽象语法树）------------------------------
    result = ''
    # 解析python代码为抽象语法树，将字符串转换为byte，并将byte流输入到解析器，生成抽象语法树
    tree = python_parser.parse(bytes(code, "utf8"))
    # ----------------------进行符号执行-----------------------------------------
    # tree.root_node抽象语法树的根节点
    root_node = tree.root_node
    # If_node是自定义的类（在entity.py）中，表示if语句节点。下面是在构建if语法树的根节点
    if_tree = If_node('module', None)
    # 下面方法从抽象语法树根节点提取if语句相关结构，并将信息存到if语法树。
    # {}是空字典 作用是在符号执行遍历if语法树时更新符号变量对照表
    format_tree(root_node, if_tree, {})
    # format_tree_cpp(root_node, if_tree, {})
    # ------------------------生成约束条件----------------------------------------
    # 打印if树
    print_tree(if_tree)
    # 生成路径约束，一条路径的约束条件是一个数组
    restrictions = get_truth_trees(if_tree)
    # 保存 if 树的图像
    images_dir = 'static/images'
    save_ast_image(if_tree, images_dir)


# 按钮-生成c++语法树
def generateTree_cpp(code: str):
    # ----------------------生成中间表示（抽象语法树）------------------------------
    result = ''
    # 解析python代码为抽象语法树，将字符串转换为byte，并将byte流输入到解析器，生成抽象语法树
    tree = cpp_parser.parse(bytes(code, "utf8"))
    # ----------------------进行符号执行-----------------------------------------
    # tree.root_node抽象语法树的根节点
    root_node = tree.root_node
    # If_node是自定义的类（在entity.py）中，表示if语句节点。下面是在构建if语法树的根节点
    if_tree = If_node('module', None)
    # 下面方法从抽象语法树根节点提取if语句相关结构，并将信息存到if语法树。
    # {}是空字典 作用是在符号执行遍历if语法树时更新符号变量对照表
    format_tree_cpp(root_node, if_tree, {})
    # ------------------------生成约束条件----------------------------------------
    # 打印if树
    # print_tree(if_tree)
    print_tree_cpp(if_tree)
    # 生成路径约束，一条路径的约束条件是一个数组
    restrictions = get_truth_trees(if_tree)
    # 保存 if 树的图像
    images_dir = 'static/images'
    save_ast_image_cpp(if_tree, images_dir)



# 按钮-生成java语法树
def generateTree_java(code: str):
    # ----------------------生成中间表示（抽象语法树）------------------------------
    result = ''
    # 解析python代码为抽象语法树，将字符串转换为byte，并将byte流输入到解析器，生成抽象语法树
    tree = java_parser.parse(bytes(code, "utf8"))
    # ----------------------进行符号执行-----------------------------------------
    # tree.root_node抽象语法树的根节点
    root_node = tree.root_node
    # If_node是自定义的类（在entity.py）中，表示if语句节点。下面是在构建if语法树的根节点
    if_tree = If_node('module', None)
    # 下面方法从抽象语法树根节点提取if语句相关结构，并将信息存到if语法树。
    # {}是空字典 作用是在符号执行遍历if语法树时更新符号变量对照表
    format_java_tree(root_node, if_tree, {})
    # ------------------------生成约束条件----------------------------------------
    # 打印if树
    # print_tree(if_tree)
    print_tree_cpp(if_tree)
    # 生成路径约束，一条路径的约束条件是一个数组
    restrictions = get_truth_trees(if_tree)
    # 保存 if 树的图像
    images_dir = 'static/images'
    save_ast_image_java(if_tree, images_dir)
